import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import os

"""
ndarray图片处理示例

这个脚本展示了如何使用numpy创建和处理多维数组(ndarray)，并使用matplotlib将其可视化为图片。

功能包括：
1. 创建不同类型的ndarray图像（随机噪声、渐变、正弦波等）
2. 对图像进行简单处理（反转、阈值化、模糊）
3. 显示处理前后的图像对比
"""

def image_to_ndarray(image_path):
    """
    从图片文件生成numpy数组

    参数:
        image_path: 图片文件路径

    返回:
        numpy数组表示的图像
    """
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"图片文件不存在: {image_path}")

    # 使用matplotlib读取图片并转换为numpy数组
    img_array = plt.imread(image_path)
    return img_array

def create_sample_images():
    """创建示例ndarray图像"""
    # 设置随机种子，确保结果可复现
    np.random.seed(42)
    
    # 1. 随机噪声图像
    noise_img = np.random.rand(256, 256)
    
    # 2. 渐变图像
    x = np.linspace(0, 1, 256)
    y = np.linspace(0, 1, 256)
    X, Y = np.meshgrid(x, y)
    gradient_img = X + Y
    
    # 3. 正弦波图案
    sine_img = np.sin(X * 10) * np.cos(Y * 10) * 0.5 + 0.5
    
    # 4. 彩色图像 (RGB)
    rgb_img = np.zeros((256, 256, 3))
    rgb_img[:, :, 0] = X  # 红色通道
    rgb_img[:, :, 1] = Y  # 绿色通道
    rgb_img[:, :, 2] = 0.5  # 蓝色通道
    
    return noise_img, gradient_img, sine_img, rgb_img

def process_image(img):
    """对图像进行简单处理"""
    # 1. 反转图像
    inverted = 1 - img if img.ndim == 2 else 1 - img
    
    # 2. 阈值化处理
    threshold = 0.5
    thresholded = (img > threshold).astype(np.float32)
    
    # 3. 简单模糊（均值滤波）
    if img.ndim == 2:
        blurred = np.zeros_like(img)
        for i in range(1, img.shape[0]-1):
            for j in range(1, img.shape[1]-1):
                blurred[i, j] = np.mean(img[i-1:i+2, j-1:j+2])
    else:
        blurred = np.zeros_like(img)
        for c in range(3):
            for i in range(1, img.shape[0]-1):
                for j in range(1, img.shape[1]-1):
                    blurred[i, j, c] = np.mean(img[i-1:i+2, j-1:j+2, c])
    
    return inverted, thresholded, blurred

def display_images(original, processed, titles):
    """显示图像及其处理结果"""
    # 创建一个包含原始图像和处理后图像的图形
    fig, axes = plt.subplots(1, len(processed) + 1, figsize=(15, 5))
    
    # 显示原始图像
    if original.ndim == 2:
        axes[0].imshow(original, cmap='gray')
    else:
        axes[0].imshow(original)
    axes[0].set_title('原始图像')
    axes[0].axis('off')
    
    # 显示处理后的图像
    for i, (img, title) in enumerate(zip(processed, titles), 1):
        if img.ndim == 2:
            axes[i].imshow(img, cmap='gray')
        else:
            axes[i].imshow(img)
        axes[i].set_title(title)
        axes[i].axis('off')
    
    plt.tight_layout()
    plt.show()


    print("处理正弦波图像...")
    inverted_sine, thresholded_sine, blurred_sine = process_image(sine_img)
    display_images(sine_img, [inverted_sine, thresholded_sine], 
                  ['反转正弦波', '阈值化正弦波'])
    
    # 处理彩色图像
    print("处理彩色图像...")
    inverted_rgb, thresholded_rgb, blurred_rgb = process_image(rgb_img)
    display_images(rgb_img, [inverted_rgb, thresholded_rgb, blurred_rgb], 
                  ['反转彩色', '阈值化彩色', '模糊彩色'])
    
    print("演示完成!")

if __name__ == "__main__":
    main()